import logging
from tqdm import tqdm
import argparse






#def cutoff(prob_0,prob_1,prob_cf):
#    if abs(prob_0-prob_1)<prob_cf:
#        return False
#    return True


def read_modkit_output_read_site(inputfile,outputfile,logger):
    i=0
    invalid_ref=0
    valid_pos=0
    infer_num=0
    outputfh = open(outputfile, "w", buffering=512)
    with open(inputfile) as inputfh:
        for line in inputfh:
            if i == 0:
                i+=1
                continue
            fields = line.split()
            ref_pos=fields[2]
            mod_base=fields[11]
            if ref_pos=='NA' or ref_pos is None:
                invalid_ref+=1
                continue
            if mod_base == 'h':
                continue
            infer=fields[-2]
            if infer=='true':
                infer_num+=1
                continue
            read_name=fields[0]
            read_pos=fields[1]
            chr=fields[3]
            
            strand=fields[5]
            read_length=fields[9]
            prob_1=float(fields[10])
            prob_0=1-prob_1
            label = 1 if prob_1 > 0.5 else 0
            kmer=fields[14]
            i+=1
            outputfh.write(
                    f"{chr}\t{ref_pos}\t{strand}\t{read_pos}\t{read_name}\t{read_length}\t"
                    f"{prob_0}\t{prob_1}\t{label}\t{kmer}\n"
                    )
    logger.info(f'invalid_ref: {invalid_ref}')
    logger.info(f'infer_num: {infer_num}')
 

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input",
        type=str,
        default="/public/home/hpc234701005/data/R10.4/20221109_1654_5A_PAG65784_f306681d/test1/deepmod2.per_read",
        help="input address.",
    )
    parser.add_argument(
        "--result",
        type=str,
        default="/public/home/hpc234701005/data/R10.4/20221109_1654_5A_PAG65784_f306681d/test1/deepmod2_methylated.tsv",
        help="result address.",
    )
    parser.add_argument(
        "--log",
        type=str,
        default="/public/home/hpc234701005/data/R10.4/log/extract_bam_by_read_site.deepmod2.log",
        help="log address.",
    )
    return parser.parse_args()

if __name__ == '__main__':
    args=parse_args()
    logger = logging.getLogger("test_logger")
    logger.setLevel(logging.DEBUG)
    read_site_log = logging.FileHandler(
        args.log, "a", encoding="utf-8"
        )
    formatter = logging.Formatter(
        "%(asctime)s - line:%(lineno)d - %(levelname)s - %(message)s"
        )
    read_site_log.setFormatter(formatter)
    # 加载文件到logger对象中
    logger.addHandler(read_site_log)
    read_modkit_output_read_site(args.input,args.result,logger)
    logger.info('----------------finish----------------------')

